from scipy.io import wavfile
import numpy as np
import wave
import matplotlib.pyplot as plt
import math

from function_test import windowing_filtered_frame,filtered_frame, recons_filtered_signal,recons_windowing_filtered_signal

class MeanVariance(object):
    def __init__(self):
        self.m2 = self.avg = 0.0
        self.maxnum = self.minnum = None
        self.num = 0

    def clear(self):
        self.m2 = self.avg = 0.0
        self.minnum = self.maxnum = None
        self.num = 0

    def reset(self, data=None):
        self.clear()
        if not data is None:
            self.add(data)

    def add(self, val):
        self.num += 1
        delta = val - self.avg
        self.avg += delta/self.num
        self.m2 += delta*(val-self.avg)
        if self.maxnum is None or val > self.maxnum:
            self.maxnum = val
        if self.minnum is None or val < self.minnum:
            self.minnum = val

    def variance(self):
        if self.num < 2:
            return 0
        return self.m2/(self.num-1)

    def std(self):
        return math.sqrt(math.fabs(self.variance()))

    def mean(self):
        return self.avg
class MeanVarianceWindow(MeanVariance):
    def __init__(self, wsize):
        super(MeanVarianceWindow, self).__init__()
        self.wsize = wsize
        self.pos = self.num = 0
        self.data = np.zeros(wsize)

    def clear(self):
        mn = self.mean()
        self.pos = self.num = 0
        self.data = np.array([mn, ]*self.wsize, dtype=float)

    def add(self, val):
        self.num += 1
        oavg = self.avg
        xold = self.data[self.pos]

        # reaches end?
        if self.num > self.wsize:
            self.num = self.wsize
            self.avg += (val-xold)/self.num
            self.m2 += (val+xold-oavg-self.avg)*(val-xold)
        else:
            self.avg += (val - self.avg)/self.num
            self.m2 += (val - oavg)*(val - self.avg)
        self.data[self.pos] = val
        self.pos += 1
        if self.pos >= self.wsize:
            self.pos = 0
        if self.maxnum is None or val > self.maxnum:
            self.maxnum = val
        if self.minnum is None or val < self.minnum:
            self.minnum = val

    def median(self):
        return np.median(np.array(self.data))

def pre_emph(signal_,pre_emphasis_):
    # start = timeit.default_timer()
    # print(timeit.default_timer()-start)
    # append `signal[-1]` is faster
    # from Bai
    return np.append(signal_[1] - pre_emphasis_ * signal_[:-1],signal_[-1])

def frame(size,stride,frequency,sig_length,input_sig,with_window):
    fra_length, fra_step = round(size * frequency), round(stride * frequency)
    fra_num = int(np.ceil(float(np.abs(sig_length - fra_length)) / fra_step))
    # creat a signal with fixed-length frames
    pad_signal_len = fra_length + fra_step * fra_num
    pad_signal_ = np.append(input_sig, np.zeros(pad_signal_len - sig_length))
    # creat a array with `num_frames` rows and `frame_length` columns
    fra = np.zeros((fra_num, fra_length))
    for n in range(0, fra_num):
        fra[n] = np.array(pad_signal_[n * fra_step:n * fra_step + fra_length])
        # frames[n] = np.array(signal[n * frame_step:n * frame_step + frame_length])
    if with_window:
        fra_window = fra * np.hamming(fra_length)
        return fra_length, fra_step, fra_num, fra,fra_window,pad_signal_
    return fra_length, fra_step, fra_num, fra,pad_signal_

def short_term_energy(frames_windowing_,num_frames_,mode = None):
    # define a function to calculate short term energy in a frame
    if mode == None or mode == 1:
        energy_abs_ = np.zeros(num_frames_)
        for i in range(0, num_frames):
            energy_abs_[i] = np.sum(np.abs(frames_windowing_[i, :]))
        return energy_abs_
    elif mode == 0:
        energy_squ_ = np.zeros(num_frames_)
        for i in range(0, num_frames):
            energy_squ_[i] = np.sum(np.square(frames_windowing_[i, :]))
        return energy_squ_
    else:
        print("error when using function 'short_term_energy'")

def zero_crossing_rate(frames_windowing_,num_frames_,delta_):
    # define a function to calculate zero crossing rate in a frame
    if num_frames_ <= 0:
        # make sure there is no error when input has one frame
        print("ZCR calculation failed!")
        return False
    else:
        zcr_ = np.zeros(num_frames_)
        for i in range(0, num_frames_):
            if num_frames_ == 1:
                single_frame = frames_windowing_
            else:
                single_frame = frames_windowing_[i]
            temp1 = single_frame[:-1]
            temp2 = single_frame[1:]
            signs = temp1 * temp2 < 0  # signs = True || False
            diff = np.abs((temp1 - temp2)) > delta_  # diff = True || False
            # zcr[i] = numpy.sum(signs * diff) / (num_frames - 1)
            zcr_[i] = np.sum(signs * diff) # logic calculation
        return zcr_

def hyst(x, th_lo, th_hi=None, initial = False):
    if th_hi is None:
        th_hi = th_lo + np.finfo(np.float32).eps
    hi = x >= th_hi
    lo_or_hi = (x <= th_lo) | hi
    ind = np.nonzero(lo_or_hi)[0]
    if not ind.size: # prevent index error if ind is empty
        return np.zeros_like(x, dtype=bool) | initial
    cnt = np.cumsum(lo_or_hi) # from 0 to len(x)
    hst = np.where(cnt, hi[ind[cnt-1]], initial)
    A, B = rising_falling_edge_det(hst)
    return hst, A, B

def rising_falling_edge_det(input):
    hst = np.array(input)
    loc = np.where(np.diff(hst))[0]
    A = loc[np.where(hst[loc] ^ True)]
    B = loc[np.where(hst[loc] ^ False)]
    if len(A) == 0 and len(B) == 0:
        print("NO rising or falling edge")
    elif len(A) == 1 and len(B) == 0:
        B = np.pad(B, (0, 1), constant_values=num_frames - 1)
    elif len(B) == 1 and len(A) == 0:
        A = np.pad(A, (0, 1), constant_values=0)
    else:
        if B[0] < A[0]:
            A = np.pad(A, (1, 0), constant_values=0)
        if len(A) != len(B):
            B = np.pad(B, (0, 1), constant_values=A[-1])
    # print("A: ", A)
    # print("B: ", B)
    return A, B

def frame_conv_point(A,B,fra_step,fra_length):
    A = A * fra_step
    B = B * fra_step + fra_length
    vad = np.zeros(num_points,dtype=bool)
    for i in range(int(len(A))):
        vad[A[i]:B[i]] = True
    return vad

def test():
    # A set of unit tests
    #assert len(energy_abs) == num_frames
    assert len(frames[0,:]) == frame_length
    assert len(frames[:,0]) == num_frames
    assert len(frames_windowing[0, :]) == frame_length
    assert len(frames_windowing[:,0]) == num_frames
    print("All test pass.")
# ---------------------------------------------------------
# read file & framing & windowing
# ---------------------------------------------------------

filename = '20200628.wav'
#filename = 'Alarm01.wav'
filename1 = 'test1.wav'
#filename = 'with_noise.wav'

sample_rate, signal = wavfile.read(filename1)
# get the basic information of this file
num_points, num_channels = signal.shape
print("{} has {} samples {} channels at {} sample rate".format(
    filename,num_points,num_channels,sample_rate
))
if num_channels != 0:
    # choose the first channel of the signal
    signal = signal[:,0]

pre_emphasis = 0.97
emphasized_signal = pre_emph(signal,pre_emphasis)

# for 25ms period, it contains 1200 points (0.025 * sample_rate = 0.025 * 48000 = 1200)
# 10ms stride (480 points) & 15ms overlap (720 points)
frame_size = 0.024      # 25ms
frame_stride = 0.01     # 10ms
signal_length = num_points
frame_length, frame_step, num_frames, frames, frames_windowing, pad_signal = \
    frame(frame_size,frame_stride,sample_rate,num_points,emphasized_signal,True)
print("frame_length: ",frame_length,"\nframe_step: ",frame_step,"\nnum_frames: ",num_frames)


test()

# ---------------------------------------------------------
# WAY 1
# using [ running average of 30 frames' energy ] and [ mean of dif std ] as threshold
# moving_ave_energy & std
# ---------------------------------------------------------
q = MeanVarianceWindow(30)
mvwofsd = MeanVarianceWindow(30)
meanofstd = MeanVarianceWindow(30)

vad_f = []
energy = 0.0
std = 0
t = []
s = []
ms = []
th = 3000
for i in range(num_frames):
    t.append(th)
    energy = np.sum(abs(frames_windowing[i,:]))
    #std = np.std(abs(frames_windowing[i,:]))
    q.add(energy)
    moving_ave_energy = q.mean()
    mvwofsd.add(energy - moving_ave_energy)
    s.append(mvwofsd.std())
    meanofstd.add(mvwofsd.std())
    std = meanofstd.mean()
    ms.append(std)
    # th = moving_ave_energy + 0*b
    if std == 0:
        vad_f.append(False)
    elif mvwofsd.std() > std/2 and energy > 400000:
        vad_f.append(True)
    else:
        #energy = np.sum(abs(frames_windowing[i,:]))
        if energy > th:
            vad_f.append(True)
        else:
            vad_f.append(False)
    if std == 0:
        th = th
    else:
        if i > 5 and all(vad_f[i-6:i]) == True:
            #th = th
            if i > 30 and all(vad_f[i-31:i]) == True:
                    th = moving_ave_energy + 0.5 * std
            else:
                th = th
        else:
            th = moving_ave_energy + 0.5 * std

# Compare threshold (t), mean of std (ms), std, frame_energy
# plt.plot(t)
# plt.plot(ms)
# plt.plot(s)
# plt.plot(short_term_energy(frames_windowing,num_frames,1))
# plt.legend()
# plt.show()
print("length of vad_f",len(vad_f))
A, B = rising_falling_edge_det(vad_f)
print("A: ", A)
print("B: ", B)
vad_p = frame_conv_point(A,B,frame_step,frame_length)

delta = 0
th_lo = 50
frame_length, frame_step, num_frames, frames, frames_windowing, pad_signal = \
    frame(frame_size,frame_stride,sample_rate,num_points,signal,True)



zcr = zero_crossing_rate(frames,num_frames,delta)
zcr_window_filtered = zero_crossing_rate(windowing_filtered_frame, len(windowing_filtered_frame), delta)
zcr_filtered = zero_crossing_rate(filtered_frame,len(filtered_frame),delta)


hi = np.array(vad_f,dtype=int)
b = (zcr <= th_lo)
lo_or_hi = b | hi
ind = np.nonzero(lo_or_hi)[0]
cnt = np.cumsum(lo_or_hi) # from 0 to len(x)
hst = np.where(cnt, hi[ind[cnt-1]], False)
C, D = rising_falling_edge_det(hst)




print("C:",C)
print("D:",D)
vad_p_zcr = frame_conv_point(C,D,frame_step,frame_length)
# ---------------------------------------------------------
# plot vad
# ---------------------------------------------------------

plt.plot(np.linspace(0, signal_length / sample_rate, signal_length), signal)
plt.plot(np.linspace(0, signal_length / sample_rate, signal_length), 20000 * vad_p, label='vad')
plt.plot(np.linspace(0, signal_length / sample_rate, signal_length), 10000 * vad_p_zcr, label='vad_ZCR')
plt.xlabel("Time[min]")
plt.ylabel("Amplitude")
plt.title("VAD for row signal using short term energy")
plt.legend()
plt.show()

# ---------------------------------------------------------
# save active signal into wave file
# ---------------------------------------------------------
result = signal * vad_p
result1 = signal * vad_p_zcr
plt.subplot(211)
plt.plot(signal,label= "original signal")
plt.legend()
plt.subplot(212)
plt.plot(result1,label = "filtered signal")
plt.legend()
plt.show()

wf_name = "vad_" + filename1
f = wave.open(wf_name,"wb")
f.setnchannels(2)
f.setsampwidth(2)
f.setframerate(sample_rate)
f.writeframes(result.tobytes())
f.close()